"""
用于模拟数据集的评分器
基于真实的依赖图（ground truth graph）进行评分
"""
from typing import Dict
from typing import Sequence
from typing import Tuple

import numpy as np

from ...alg.ci import RHTScorer
from ...model.case import CaseData
from ...model.graph import Node


class SimRHTScorer(RHTScorer):
    """
    模拟数据集专用的回归假设检验评分器
    
    扩展了基础RHTScorer，使用真实的父节点关系：
    Pa(V_{i}^{(t)}) = Pa^{(t)}(V_{i}^{t}) \\cup {V_{i}^{(t - 1)}}
    
    其中：
    - V_{i}^{(t)} 表示节点i在时间t的状态
    - Pa^{(t)}(V_{i}^{t}) 表示节点i在时间t的父节点集合
    - V_{i}^{(t - 1)} 表示节点i在时间t-1的状态（自身的历史状态）
    """

    def split_data(
        self,
        data: Dict[Node, Sequence[float]],  # 所有节点的时间序列数据
        node: Node,                         # 目标节点
        parents: Sequence[Node],            # 父节点列表
        case_data: CaseData,               # 案例数据
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        """
        准备训练和测试数据
        
        特点：
        1. 考虑节点的父节点影响
        2. 加入节点自身的历史状态（t-1时刻）
        3. 处理时间序列的对齐问题
        
        Returns:
            (train_x, test_x, train_y, test_y): 训练集和测试集数据
        """
        # 收集所有父节点的时间序列
        series = [data[parent] for parent in parents if parent in data]
        
        # 添加节点自身的历史状态（将序列向后移动一位）
        # 例如：[x1, x2, ..., xn] 变为 [xn, x1, x2, ...]
        series.append(np.roll(data[node], 1))
        
        # 转置数据矩阵，使每行代表一个时间点的特征
        series = np.array(series).T
        
        # 去除第一个时间点（因为它没有历史状态）
        series_x: np.ndarray = series[1:, :]  # 特征矩阵
        series_y = np.array(data[node][1:])   # 目标值序列

        # 分割训练集和测试集
        return self._split_train_test(
            series_x=series_x,
            series_y=series_y,
            train_window=case_data.train_window - 1,  # 减1是因为去掉了第一个时间点
            test_window=case_data.test_window,
        )
